"""
Function test without fastapi
"""
import os
import time
from engine import SearchEngine
from engine.utils import get_config, vis
from engine.models.resnet import ResnetFeatureExtractor
from engine.utils.mysql_helper import MySQL
from engine.models import BiTFeatureExtractor, MAEFeatureExtractor


def test_resnet():
    model = ResnetFeatureExtractor()
    image_path = 'data/000048.jpg'
    feat = model.extract(image_path)
    print(feat.shape)


def test_bit():
    bit_name = 'BiT-M-R50x1'
    extractor = BiTFeatureExtractor(bit_name)
    image_path = './data/000048.jpg'
    feature = extractor.extract(image_path)
    print(feature.shape)


def test_mae():
    bit_name = 'vit_base_patch16'
    extractor = MAEFeatureExtractor(bit_name)
    image_path = './data/000048.jpg'
    feature = extractor.extract(image_path)
    print(feature.shape)


def test_engine_insert():
    config_path = 'cfg/resnet.yaml'
    config = get_config(config_path)
    engine = SearchEngine(config)

    table_name = config.TABLE_NAME
    image_folder = '/rinoshinme/data/image_search'

    # clear data and insert elements
    engine.clear(table_name)
    print('data cleared')

    # start upload
    for root, _, files in os.walk(image_folder):
        for name in files:
            if not name.endswith('.jpg'):
                continue
            image_path = os.path.join(root, name)
            print('uploading {}'.format(image_path))
            ts = int(time.time())
            engine.upload(table_name, image_path, ts)


def test_engine_search():
    config_path = 'cfg/resnet.yaml'
    config = get_config(config_path)
    engine = SearchEngine(config)

    table_name = config.TABLE_NAME
    image_folder = '/rinoshinme/data/image_search'

    # search
    nmatch = 5
    names = os.listdir(image_folder)
    query_image = os.path.join(image_folder, names[0])
    result_paths, distances, ids = engine.search(table_name, query_image, top_k=nmatch)
    print(result_paths)
    print(distances)
    print(ids)
    # vis(query_image, result_paths, distances)


def test_engine_count():
    config_path = 'cfg/resnet.yaml'
    config = get_config(config_path)
    engine = SearchEngine(config)

    table_name = config.TABLE_NAME
    print(engine.get_database_size(table_name))


def test_mysql():
    config_path = 'cfg/resnet.yaml'
    config = get_config(config_path)
    client = MySQL(config)
    client.test_connection()


def main():
    test_resnet()
    test_bit()
    test_mae()
    # test_engine_insert()
    # test_engine_search()
    # test_engine_count()


if __name__ == '__main__':
    main()
